package main

import (
    "ebbflow/pkg/crypto"
    "encoding/json"
    "flag"
    "fmt"
    "log"
    "os"
    "path/filepath"
)

// BLSKeyConfig represents the configuration for BLS key generation
type BLSKeyConfig struct {
    N int `json:"n"` // Total number of nodes
    T int `json:"t"` // Threshold (minimum number of shares needed)
}

// BLSKeySet represents a complete set of BLS keys for all nodes
type BLSKeySet struct {
    Config     BLSKeyConfig     `json:"config"`
    ThresholdPK crypto.G1       `json:"threshold_pk"`
    PublicKeys []crypto.G1      `json:"public_keys"`
    PrivateKeys []crypto.ZR     `json:"private_keys"`
}

// NodeKeyFile represents the key file for a single node
type NodeKeyFile struct {
    NodeID     int         `json:"node_id"`
    Config     BLSKeyConfig `json:"config"`
    ThresholdPK crypto.G1   `json:"threshold_pk"`
    PublicKeys []crypto.G1  `json:"public_keys"`
    PrivateKey crypto.ZR    `json:"private_key"`
}

// BLSKeyGenerator handles BLS key generation and file management
type BLSKeyGenerator struct {
    outputDir string
    logger    *log.Logger
}

func NewBLSKeyGenerator(outputDir string) *BLSKeyGenerator {
    return &BLSKeyGenerator{
        outputDir: outputDir,
        logger:    log.New(os.Stdout, "[BLS-GEN] ", log.LstdFlags),
    }
}

// GenerateKeys generates threshold BLS keys for n nodes with threshold t
func (g *BLSKeyGenerator) GenerateKeys(n, t int) (*BLSKeySet, error) {
    g.logger.Printf("Generating BLS keys: n=%d, t=%d", n, t)
    
    if t > n {
        return nil, fmt.Errorf("threshold t=%d cannot be greater than n=%d", t, n)
    }
    
    if t < 1 {
        return nil, fmt.Errorf("threshold t=%d must be at least 1", t)
    }
    
    // Generate threshold keys
    thpk, _, thsks := crypto.GenerateThresholdKeys(n, t)
    
    // Generate public keys
    g1 := crypto.GetG1()
    pks := make([]crypto.G1, n)
    for i := 0; i < n; i++ {
        pks[i] = crypto.MulG1(&g1, &thsks[i])
    }
    
    keySet := &BLSKeySet{
        Config: BLSKeyConfig{N: n, T: t},
        ThresholdPK: *thpk,
        PublicKeys: pks,
        PrivateKeys: thsks,
    }
    
    g.logger.Printf("Generated BLS keys successfully")
    return keySet, nil
}

// SaveKeySet saves the complete key set to files
func (g *BLSKeyGenerator) SaveKeySet(keySet *BLSKeySet) error {
    // Create output directory
    if err := os.MkdirAll(g.outputDir, 0755); err != nil {
        return fmt.Errorf("failed to create output directory: %v", err)
    }
    
    // Save complete key set
    completeFile := filepath.Join(g.outputDir, "bls_keys_complete.json")
    if err := g.saveJSON(completeFile, keySet); err != nil {
        return fmt.Errorf("failed to save complete key set: %v", err)
    }
    
    // Save individual node key files
    for i := 0; i < keySet.Config.N; i++ {
        nodeFile := &NodeKeyFile{
            NodeID: i,
            Config: keySet.Config,
            ThresholdPK: keySet.ThresholdPK,
            PublicKeys: keySet.PublicKeys,
            PrivateKey: keySet.PrivateKeys[i],
        }
        
        filename := filepath.Join(g.outputDir, fmt.Sprintf("bls_keys_node_%d.json", i))
        if err := g.saveJSON(filename, nodeFile); err != nil {
            return fmt.Errorf("failed to save node %d key file: %v", i, err)
        }
    }
    
    // Save configuration summary
    configFile := filepath.Join(g.outputDir, "bls_config.json")
    if err := g.saveJSON(configFile, keySet.Config); err != nil {
        return fmt.Errorf("failed to save config file: %v", err)
    }
    
    g.logger.Printf("Saved key files to directory: %s", g.outputDir)
    return nil
}

// saveJSON saves a struct to a JSON file
func (g *BLSKeyGenerator) saveJSON(filename string, data interface{}) error {
    file, err := os.Create(filename)
    if err != nil {
        return err
    }
    defer file.Close()
    
    encoder := json.NewEncoder(file)
    encoder.SetIndent("", "  ")
    return encoder.Encode(data)
}

// ValidateKeySet validates a generated key set
func (g *BLSKeyGenerator) ValidateKeySet(keySet *BLSKeySet) error {
    g.logger.Printf("Validating generated key set...")
    
    // Check configuration
    if keySet.Config.N <= 0 {
        return fmt.Errorf("invalid n: %d", keySet.Config.N)
    }
    if keySet.Config.T <= 0 || keySet.Config.T > keySet.Config.N {
        return fmt.Errorf("invalid threshold t: %d (n=%d)", keySet.Config.T, keySet.Config.N)
    }
    
    // Check key counts
    if len(keySet.PublicKeys) != keySet.Config.N {
        return fmt.Errorf("public key count mismatch: expected %d, got %d", 
            keySet.Config.N, len(keySet.PublicKeys))
    }
    if len(keySet.PrivateKeys) != keySet.Config.N {
        return fmt.Errorf("private key count mismatch: expected %d, got %d", 
            keySet.Config.N, len(keySet.PrivateKeys))
    }
    
    // Validate threshold signature capability
    testMessage := []byte("test_message_for_validation")
    
    // Test threshold signature with minimum required shares
    signerSet := make([]int, keySet.Config.T)
    sigShares := make(map[int]*crypto.G2)
    
    for i := 0; i < keySet.Config.T; i++ {
        signerSet[i] = i
        sigShares[i] = crypto.Sign(&keySet.PrivateKeys[i], testMessage)
    }
    
    // Combine signatures
    aggSig := crypto.CombineShares(sigShares, signerSet)
    
    // Verify against threshold public key (simplified verification)
    if aggSig == nil {
        return fmt.Errorf("threshold signature combination failed")
    }
    
    g.logger.Printf("Key set validation passed")
    return nil
}

// GenerateHostsTemplate generates a hosts template file
func (g *BLSKeyGenerator) GenerateHostsTemplate(n int) error {
    hostsFile := filepath.Join(g.outputDir, "hosts.template")
    
    file, err := os.Create(hostsFile)
    if err != nil {
        return fmt.Errorf("failed to create hosts template: %v", err)
    }
    defer file.Close()
    
    fmt.Fprintf(file, "# ADKR Hosts Template\n")
    fmt.Fprintf(file, "# Generated for %d nodes\n", n)
    fmt.Fprintf(file, "# Format: <node_id> <hostname_or_ip> <port>\n\n")
    
    for i := 0; i < n; i++ {
        fmt.Fprintf(file, "%d localhost %d\n", i, 8000+i)
    }
    
    g.logger.Printf("Generated hosts template: %s", hostsFile)
    return nil
}

// GenerateDeploymentScript generates a deployment script
func (g *BLSKeyGenerator) GenerateDeploymentScript(n int) error {
    scriptFile := filepath.Join(g.outputDir, "deploy.sh")
    
    file, err := os.Create(scriptFile)
    if err != nil {
        return fmt.Errorf("failed to create deployment script: %v", err)
    }
    defer file.Close()
    
    fmt.Fprintf(file, "#!/bin/bash\n")
    fmt.Fprintf(file, "# ADKR Deployment Script\n")
    fmt.Fprintf(file, "# Generated for %d nodes\n\n", n)
    
    fmt.Fprintf(file, "set -e\n\n")
    
    fmt.Fprintf(file, "BINARY_PATH=\"./adkr_node\"\n")
    fmt.Fprintf(file, "HOSTS_FILE=\"hosts.template\"\n\n")
    
    fmt.Fprintf(file, "echo \"Starting ADKR deployment...\"\n\n")
    
    for i := 0; i < n; i++ {
        fmt.Fprintf(file, "echo \"Starting node %d...\"\n", i)
        fmt.Fprintf(file, "$BINARY_PATH -pid %d -hosts $HOSTS_FILE -key bls_keys_node_%d.json &\n", i, i)
        fmt.Fprintf(file, "sleep 1\n\n")
    }
    
    fmt.Fprintf(file, "echo \"All nodes started. Use 'pkill adkr_node' to stop.\"\n")
    
    // Make script executable
    if err := os.Chmod(scriptFile, 0755); err != nil {
        return fmt.Errorf("failed to make script executable: %v", err)
    }
    
    g.logger.Printf("Generated deployment script: %s", scriptFile)
    return nil
}

func main() {
    var (
        n         = flag.Int("n", 4, "Number of nodes")
        t         = flag.Int("t", 3, "Threshold (minimum shares needed)")
        outputDir = flag.String("output", "./keys", "Output directory for key files")
        validate  = flag.Bool("validate", true, "Validate generated keys")
        hosts     = flag.Bool("hosts", true, "Generate hosts template")
        deploy    = flag.Bool("deploy", true, "Generate deployment script")
    )
    flag.Parse()
    
    log.Printf("BLS Key Generator - n=%d, t=%d, output=%s", *n, *t, *outputDir)
    
    generator := NewBLSKeyGenerator(*outputDir)
    
    // Generate keys
    keySet, err := generator.GenerateKeys(*n, *t)
    if err != nil {
        log.Fatalf("Failed to generate keys: %v", err)
    }
    
    // Validate keys if requested
    if *validate {
        if err := generator.ValidateKeySet(keySet); err != nil {
            log.Fatalf("Key validation failed: %v", err)
        }
    }
    
    // Save key files
    if err := generator.SaveKeySet(keySet); err != nil {
        log.Fatalf("Failed to save key files: %v", err)
    }
    
    // Generate hosts template if requested
    if *hosts {
        if err := generator.GenerateHostsTemplate(*n); err != nil {
            log.Fatalf("Failed to generate hosts template: %v", err)
        }
    }
    
    // Generate deployment script if requested
    if *deploy {
        if err := generator.GenerateDeploymentScript(*n); err != nil {
            log.Fatalf("Failed to generate deployment script: %v", err)
        }
    }
    
    log.Printf("BLS key generation completed successfully!")
    log.Printf("Files generated in: %s", *outputDir)
    log.Printf("- bls_keys_complete.json: Complete key set")
    log.Printf("- bls_keys_node_*.json: Individual node key files")
    log.Printf("- bls_config.json: Configuration summary")
    if *hosts {
        log.Printf("- hosts.template: Hosts configuration template")
    }
    if *deploy {
        log.Printf("- deploy.sh: Deployment script")
    }
}